{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "32Fac6LDZJmE",
    "colab_type": "code",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 125.0
    },
    "outputId": "7c9fab8d-0e3f-4994-aecc-0b3d00a59ffc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1000 cost = 0.064772\n",
      "Epoch: 2000 cost = 0.014707\n",
      "Epoch: 3000 cost = 0.005852\n",
      "Epoch: 4000 cost = 0.002837\n",
      "Epoch: 5000 cost = 0.001509\n",
      "[['i', 'like'], ['i', 'love'], ['i', 'hate']] -> ['dog', 'coffee', 'milk']\n"
     ]
    }
   ],
   "source": [
    "'''\n",
    "  code by Tae Hwan Jung(Jeff Jung) @graykode\n",
    "'''\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "\n",
    "dtype = torch.FloatTensor\n",
    "\n",
    "sentences = [ \"i like dog\", \"i love coffee\", \"i hate milk\"]\n",
    "\n",
    "word_list = \" \".join(sentences).split()\n",
    "word_list = list(set(word_list))\n",
    "word_dict = {w: i for i, w in enumerate(word_list)}\n",
    "number_dict = {i: w for i, w in enumerate(word_list)}\n",
    "n_class = len(word_dict)\n",
    "\n",
    "# TextRNN Parameter\n",
    "batch_size = len(sentences)\n",
    "n_step = 2 # number of cells(= number of Step)\n",
    "n_hidden = 5 # number of hidden units in one cell\n",
    "\n",
    "def make_batch(sentences):\n",
    "    input_batch = []\n",
    "    target_batch = []\n",
    "\n",
    "    for sen in sentences:\n",
    "        word = sen.split()\n",
    "        input = [word_dict[n] for n in word[:-1]]\n",
    "        target = word_dict[word[-1]]\n",
    "\n",
    "        input_batch.append(np.eye(n_class)[input])\n",
    "        target_batch.append(target)\n",
    "\n",
    "    return input_batch, target_batch\n",
    "\n",
    "# to Torch.Tensor\n",
    "input_batch, target_batch = make_batch(sentences)\n",
    "input_batch = Variable(torch.Tensor(input_batch))\n",
    "target_batch = Variable(torch.LongTensor(target_batch))\n",
    "\n",
    "class TextRNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(TextRNN, self).__init__()\n",
    "\n",
    "        self.rnn = nn.RNN(input_size=n_class, hidden_size=n_hidden)\n",
    "        self.W = nn.Parameter(torch.randn([n_hidden, n_class]).type(dtype))\n",
    "        self.b = nn.Parameter(torch.randn([n_class]).type(dtype))\n",
    "\n",
    "    def forward(self, hidden, X):\n",
    "        X = X.transpose(0, 1) # X : [n_step, batch_size, n_class]\n",
    "        outputs, hidden = self.rnn(X, hidden)\n",
    "        # outputs : [n_step, batch_size, num_directions(=1) * n_hidden]\n",
    "        # hidden : [num_layers(=1) * num_directions(=1), batch_size, n_hidden]\n",
    "        outputs = outputs[-1] # [batch_size, num_directions(=1) * n_hidden]\n",
    "        model = torch.mm(outputs, self.W) + self.b # model : [batch_size, n_class]\n",
    "        return model\n",
    "\n",
    "model = TextRNN()\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# Training\n",
    "for epoch in range(5000):\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # hidden : [num_layers * num_directions, batch, hidden_size]\n",
    "    hidden = Variable(torch.zeros(1, batch_size, n_hidden))\n",
    "    # input_batch : [batch_size, n_step, n_class]\n",
    "    output = model(hidden, input_batch)\n",
    "\n",
    "    # output : [batch_size, n_class], target_batch : [batch_size] (LongTensor, not one-hot)\n",
    "    loss = criterion(output, target_batch)\n",
    "    if (epoch + 1) % 1000 == 0:\n",
    "        print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))\n",
    "\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "input = [sen.split()[:2] for sen in sentences]\n",
    "\n",
    "# Predict\n",
    "hidden = Variable(torch.zeros(1, batch_size, n_hidden))\n",
    "predict = model(hidden, input_batch).data.max(1, keepdim=True)[1]\n",
    "print([sen.split()[:2] for sen in sentences], '->', [number_dict[n.item()] for n in predict.squeeze()])"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "TextRNN-Torch.ipynb",
   "version": "0.3.2",
   "provenance": [],
   "collapsed_sections": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "accelerator": "GPU"
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
